#!/usr/bin/env python
"""
Script to retrieve accessions assigned to specific taxa. Searches
through SSUMMO results files (.pkl files) and retrieves all 
accessions asigned to the taxa of interest.
"""
import cPickle as pickle
import os
import ssummolib
import re
import sys
from Bio import SeqIO

class Finder( ):
    def __init__( self, cwd='.' ):
        self.wd = cwd
    def search( self,results_dict , name ):
        reg = re.compile(name,re.I)
        for path, node in ssummolib.dict_walk( '', results_dict ):
            for OTU in node.keys():
                if reg.search(OTU):
                    sub_node = results_dict
                    for comp in path.split( os.path.sep ):
                        sub_node = sub_node[comp]
                    for path, sub_node in ssummolib.dict_walk( '', sub_node ):
                        if 'accessions' in sub_node.keys():
                            for accession in sub_node['accessions']:
                                yield accession

    def find_seq_file( self, dataset_name ):
        if '.' in dataset_name:
            prefix, suffix = dataset_name.rsplit('.',1)
        else:
            prefix = dataset_name

        # Add some common fasta suffixes.
        fasta_formats = [ 'fas','fsa' ]
        for fasta in fasta_formats:  
            SeqIO._FormatToIterator.update( { fasta : SeqIO._FormatToIterator['fasta'] } )

        for fname in os.listdir( self.wd ):
            if fname.startswith(prefix):
                seq_format = self.get_format( fname )
                if seq_format:
                    sequence_file = fname
                    return sequence_file, seq_format
                    break
        return False, False  # Shouldn't get here if finds a sequence file.

    def get_format( self, file_name ):
        suffix = file_name.rsplit('.',1)[1]
        if suffix in SeqIO._FormatToIterator.keys():
            return suffix
        else:
            return False

    def get_sequences( self,dataset_name , accessions ):
        """Dataset_name is the name of the dataset file.
        sequence format will be auto-detected.
        accessions should be an iterable item that yields
        accession numbers to search for in the file."""
        seq_file_name , seq_format = self.find_seq_file( dataset_name )
        if seq_file_name not in os.listdir(self.wd) or not seq_file_name:
            #sys.stderr.write( "Can't find {0} in {1}".format( dataset_name, self.wd) )
            return
        if seq_format in SeqIO._BinaryFormats:
            handle = file( seq_file_name , 'rb' )
        else:
            handle = file( seq_file_name , 'r' )
        accessions = set( [acc for acc in accessions] )
        for seq in SeqIO.parse( handle, seq_format ):
            if seq.id in accessions:
                accessions.remove( seq.id )
                yield seq.format( 'fasta' )
        handle.close()


if __name__ == '__main__':
    seq_finder = Finder()

    for file_name in os.listdir('.'):
        if file_name.endswith('.pkl'):
            with file(file_name,'rb') as in_file:
                result_dict = pickle.load( in_file )
                accessions = seq_finder.search( result_dict , 'Treponema' )
                for seq in seq_finder.get_sequences( file_name , accessions ):
                    sys.stdout.write( seq )

